// =============================================================================
//                             头文件引入
// =============================================================================
#include <stdlib.h>           // 标准库，用于内存分配(malloc)、字符串转换(atoi)等
#include <string.h>           // 字符串操作函数(memset, memcpy)
#include <stdio.h>            // 标准输入输出(printf)
#include <errno.h>            // 错误码定义(errno)
#include <sys/types.h>        // 基本系统数据类型
#include <netinet/in.h>       // 互联网地址族定义(sockaddr_in)
#include <arpa/inet.h>        // IP地址转换函数(inet_ntop)
#include <sys/socket.h>       // 套接字接口
#include <netdb.h>            // 网络数据库操作(getaddrinfo)
#include <byteswap.h>         // 字节序转换
#include <unistd.h>           // POSIX操作系统API(sleep, getopt)
#include <getopt.h>           // 命令行参数解析
#include <rdma/rdma_cma.h>    // RDMA CM (Communication Manager) 库的核心头文件



// =============================================================================
//                         自定义数据结构
// =============================================================================

/**
 *! @brief 代表一个 RDMA 通信节点。 它封装了与单个连接或端点相关的所有 RDMA 资源。
 */
struct cmatest_node {
    int id;                      // 节点的唯一标识符 (在此示例中未使用)
    struct rdma_cm_id *cma_id;   // RDMA CM 标识符，是 librdmacm 中最重要的对象。
                                 // 它代表一个通信端点，并包含 QP、设备上下文等信息。
    int connected;               // 一个标志，表示该节点是否已成功加入多播组并准备好通信。
    struct ibv_pd *pd;           // 保护域 (Protection Domain)。是 RDMA 资源的容器，
                                 // 如 QP、CQ、MR 都必须与一个 PD 关联。
    struct ibv_cq *cq;           // 完成队列 (Completion Queue)。用于存放已完成的工作请求(WR)
                                 // 的状态信息，即工作完成(WC)。
    struct ibv_mr *mr;           // 内存区域 (Memory Region)。一块被注册用于 RDMA 操作的内存。
                                 // 只有注册过的内存才能被 RDMA 网卡直接访问。
    struct ibv_ah *ah;           // 地址句柄 (Address Handle)。包含了路由到远程节点所需的所有
                                 // 信息 (如LID, GID)。对于 UD QP 发送数据是必需的。
    uint32_t remote_qpn;         // 远程队列对编号 (Queue Pair Number)。
    uint32_t remote_qkey;        // 远程队列密钥 (Queue Key)。用于对传入的数据包进行校验。
    void *mem;                   // 指向用户空间内存缓冲区的指针，这块内存将被注册为 MR。
};

/**
 *! @brief 代表整个测试程序的全局上下文。
 */
struct cmatest {
    struct rdma_event_channel *channel;       // RDMA 事件通道。应用程序通过此通道接收来自 librdmacm 的事件
                                              // (如地址解析完成、加入多播组成功等)。
    struct cmatest_node       *nodes;         // 指向 cmatest_node 结构体数组的指针。
    int                       conn_index;     // 连接索引 (在此示例中未使用)
    int                       connects_left;  // 剩余待完成的连接数。用于同步，等待所有节点都准备好。
    struct sockaddr_in6       dst_in;         // 存储目标多播地址 (IPv6 格式)。
    struct sockaddr           *dst_addr;      // 指向目标地址的通用指针。
    struct sockaddr_in6       src_in;         // 存储本地绑定地址 (IPv6 格式)。
    struct sockaddr           *src_addr;      // 指向本地地址的通用指针。
};



// =============================================================================
//                             全局变量
// =============================================================================
static struct cmatest       test;
static int                  connections = 1;          // 要创建的连接数，默认为1。
static int                  message_size = 100;       // 每条消息的大小 (字节)，对应 -S 参数。
static int                  message_count = 10;       // 要传输的消息数量，对应 -C 参数。
static int                  is_sender;                // 标志位，如果为 1，则程序作为发送端运行，对应 -s 参数。
static int                  unmapped_addr;            // 未使用的标志。
static char                 *dst_addr;                // 字符串形式的目标多播地址，来自 -m 参数。
static char                 *src_addr;                // 字符串形式的本地绑定地址，来自 -b 参数。
static enum rdma_port_space port_space = RDMA_PS_UDP; // RDMA 端口空间，默认为 UDP。




// =============================================================================
//                          函数实现
// =============================================================================

/**
 *! @brief 为节点分配并注册用于 RDMA 通信的内存。
 */
static int create_message(struct cmatest_node *node) {
    if (!message_size)
        message_count = 0;

    if (!message_count)
        return 0;


    // 为消息数据和 GRH (Global Route Header) 分配内存。
    // UD 类型的 QP 接收数据时，数据包前会附带 40 字节的 GRH。
    node->mem = malloc(message_size + sizeof(struct ibv_grh));
    if (!node->mem) {
        printf("failed message allocation\n");
        return -1;
    }


    // 调用 ibv_reg_mr 注册内存。这是 RDMA 的关键步骤，它会 "钉住" (pin)
    // 这块内存，防止操作系统将其交换到磁盘，并返回一个内存句柄(MR)。
    // 参数:
    //  - node->pd: 关联的保护域。
    //  - node->mem: 要注册的内存缓冲区的起始地址。
    //  - size: 要注册的内存大小。
    //  - IBV_ACCESS_LOCAL_WRITE: 访问权限，表示本地 CPU可以写入这块内存。
    node->mr = ibv_reg_mr(node->pd, 
                        node->mem, 
                        message_size + sizeof(struct ibv_grh),
                        IBV_ACCESS_LOCAL_WRITE
                    );

    if (!node->mr) {
        printf("failed to reg MR\n");
        goto err;
    }

    return 0;

err:
    free(node->mem);
    return -1;
}



/**
 *! @brief 验证测试参数是否有效，特别是消息大小不能超过网卡的 MTU。
 */
static int verify_test_params(struct cmatest_node *node) {
    struct ibv_port_attr port_attr;
    int ret;

    // 查询 RDMA 端口的属性。
    ret = ibv_query_port(node->cma_id->verbs, 
                        node->cma_id->port_num, 
                        &port_attr
                    );
    if (ret)
        return ret;


    // 检查消息大小是否超过了端口当前活动的 MTU (Maximum Transmission Unit)。
    // UD 数据包不能被分片，所以必须小于 MTU。
    if (message_count && message_size > (1 << (port_attr.active_mtu + 7))) {
        printf("mckey: message_size %d is larger than active mtu %d\n", 
                message_size, 
                1 << (port_attr.active_mtu + 7)
            );
        return -EINVAL;
    }

    return 0;
}



/**
 *! @brief 初始化一个节点的所有 RDMA 资源 (PD, CQ, QP, MR)。
 */
static int init_node(struct cmatest_node *node) {
    struct ibv_qp_init_attr init_qp_attr;
    int cqe, ret;

    // 1. 分配保护域 (PD)
    node->pd = ibv_alloc_pd(node->cma_id->verbs);
    if (!node->pd) {
        ret = -ENOMEM;
        printf("mckey: unable to allocate PD\n");
        return ret;
    }

    // 2. 创建完成队列 (CQ)
    cqe = message_count ? message_count * 2 : 2;
    node->cq = ibv_create_cq(node->cma_id->verbs, cqe, node, 0, 0);
    if (!node->cq) {
        ret = -ENOMEM;
        printf("mckey: unable to create CQ\n");
        return ret;
    }


    // 3. 准备创建队列对 (QP) 的属性
    memset(&init_qp_attr, 0, sizeof init_qp_attr);
    init_qp_attr.cap.max_send_wr = message_count ? message_count : 1; // 发送队列深度
    init_qp_attr.cap.max_recv_wr = message_count ? message_count : 1; // 接收队列深度

    init_qp_attr.cap.max_send_sge = 1;   // 发送SGE最大数量
    init_qp_attr.cap.max_recv_sge = 1;   // 接收SGE最大数量

    init_qp_attr.qp_context = node;
    init_qp_attr.sq_sig_all = 0;

    init_qp_attr.qp_type = IBV_QPT_UD;  // **关键**: QP 类型设置为不可靠数据报 (UD)，这是多播的基础。
    init_qp_attr.send_cq = node->cq;    // 关联发送完成队列
    init_qp_attr.recv_cq = node->cq;    // 关联接收完成队列


    // 4. 创建 QP
    // rdma_create_qp 是 librdmacm 提供的便捷函数，它封装了 ibv_create_qp 和状态转换。
    ret = rdma_create_qp(node->cma_id, node->pd, &init_qp_attr);
    if (ret) {
        printf("mckey: unable to create QP: %d\n", ret);
        return ret;
    }

    ret = create_message(node);
    if (ret){
        printf("mckey: failed to create messages: %d\n", ret);
    }
    return ret;
}



/**
 *! @brief 向接收队列中投递接收工作请求 (Recv WRs)。
 *        接收端必须提前准备好缓冲区来接收数据。
 */
static int post_recvs(struct cmatest_node *node) {
    struct ibv_recv_wr recv_wr, *recv_failure;
    struct ibv_sge sge;
    int i, ret = 0;

    if (!message_count)
        return 0;


    // 准备一个接收工作请求 (WR) 的模板
    recv_wr.next    = NULL;
    recv_wr.sg_list = &sge;    // 指向 Scatter/Gather Entry 列表
    recv_wr.num_sge = 1;
    recv_wr.wr_id   = (uintptr_t) node;

    // 准备 SGE，它描述了一块用于数据传输的内存缓冲区   
    sge.length      = message_size + sizeof(struct ibv_grh);  // 缓冲区长度
    sge.lkey        = node->mr->lkey;                         // 内存区域的本地密钥
    sge.addr        = (uintptr_t) node->mem;                  // 缓冲区地址


    // 循环投递 message_count 个接收请求
    for (i = 0; i < message_count && !ret; i++ ) {
        // ibv_post_recv 是一个异步调用，它将接收请求放入 QP 的接收队列。
        ret = ibv_post_recv(node->cma_id->qp, 
                            &recv_wr, 
                            &recv_failure
                        );

        if (ret) {
            printf("failed to post receives: %d\n", ret);
            break;
        }
    }
    return ret;
}



/**
 *! @brief 向发送队列中投递发送工作请求 (Send WRs)。
 */
static int post_sends(struct cmatest_node *node, 
                    int signal_flag) {
    struct ibv_send_wr send_wr, *bad_send_wr;
    struct ibv_sge sge;

    int i, ret = 0;

    if (!node->connected || !message_count)
        return 0;


    // 准备一个发送工作请求 (WR) 的模板
    send_wr.next       = NULL;
    send_wr.sg_list    = &sge;
    send_wr.num_sge    = 1;
    send_wr.opcode     = IBV_WR_SEND_WITH_IMM;    // 操作码：发送（带立即数）
    send_wr.send_flags = signal_flag;             // 发送标志，0表示非信令模式，可以提高性能


    // **关键**: 设置 UD QP 发送所需的目标地址信息
    send_wr.wr_id    = (unsigned long)node;
    send_wr.imm_data = htonl(node->cma_id->qp->qp_num);

    send_wr.wr.ud.ah          = node->ah;            // 目标地址句柄
    send_wr.wr.ud.remote_qpn  = node->remote_qpn;    // 目标 QP 编号
    send_wr.wr.ud.remote_qkey = node->remote_qkey;   // 目标 QP 密钥
    
    sge.length = message_size;
    sge.lkey   = node->mr->lkey;
    sge.addr   = (uintptr_t) node->mem;

    // 循环投递 message_count 个发送请求
    for (i = 0; i < message_count && !ret; i++) {
        ret = ibv_post_send(node->cma_id->qp, 
                            &send_wr, 
                            &bad_send_wr
                        );
        if (ret)
            printf("failed to post sends: %d\n", ret);
    }

    return ret;
}


static void connect_error(void) {
    test.connects_left--;
}


// =============================================================================
//                       librdmacm 事件处理函数
// =============================================================================



/**
 *! @brief `RDMA_CM_EVENT_ADDR_RESOLVED` 事件的处理函数。
 *        当 IP 地址成功解析为 RDMA 地址(GID)后，此函数被调用。
 */
static int addr_handler(struct cmatest_node *node) {
    int ret;

    //  1. 验证参数
    ret = verify_test_params(node);
    if (ret)
        goto err;

    // 2. 初始化所有 RDMA 资源 (PD, CQ, QP, MR)
    ret = init_node(node);
    if (ret)
        goto err;


    // 3. 如果是接收端，提前投递接收请求
    if (!is_sender) {
        ret = post_recvs(node);
        if (ret)
        goto err;
    }

    // 4. **关键**: 调用 rdma_join_multicast 加入多播组。
    //    这是一个异步操作，成功后会产生一个 RDMA_CM_EVENT_MULTICAST_JOIN 事件。
    //    参数:
    //      - node->cma_id: 当前的 CM ID。
    //      - test.dst_addr: 要加入的多播组地址。
    //      - node: 用户上下文，在事件回调中可以取回。
    ret = rdma_join_multicast(node->cma_id, 
                            test.dst_addr, 
                            node
                        );

    if (ret) {
        printf("mckey: failure joining: %d\n", ret);
        goto err;
    }
    return 0;

err:
    connect_error();
    return ret;
}



/**
 *! @brief `RDMA_CM_EVENT_MULTICAST_JOIN` 事件的处理函数。
 *        当成功加入多播组后，此函数被调用。
 */
static int join_handler(struct cmatest_node *node, 
                    struct rdma_ud_param *param) {
    char buf[40];

    inet_ntop(AF_INET6, param->ah_attr.grh.dgid.raw, buf, 40);
    printf("mckey: joined dgid: %s\n", buf);

    // 1. 保存从事件参数中获取的多播通信所需信息
    node->remote_qpn = param->qp_num;   // 实际上对于多播发送，这个值应被忽略，使用 0xFFFFFF。
    node->remote_qkey = param->qkey;    // 多播 QKey，通常是固定的。


    // 2. **关键**: 使用事件参数中的 ah_attr 创建地址句柄 (AH)。
    //    这个 AH 包含了发送多播数据包所需的所有路由信息。
    node->ah = ibv_create_ah(node->pd, &param->ah_attr);

    if (!node->ah) {
        printf("mckey: failure creating address handle\n");
        goto err;
    }

    // 3. 设置连接状态为成功，并减少等待计数
    node->connected = 1;
    test.connects_left--;
    return 0;

err:
    connect_error();
    return -1;
}



/**
 *! @brief `rdma_get_cm_event` 的主回调函数，根据事件类型分发给不同的处理函数。
 */
static int cma_handler(struct rdma_cm_id *cma_id, 
                    struct rdma_cm_event *event) {
    int ret = 0;

    switch (event->event) {
        case RDMA_CM_EVENT_ADDR_RESOLVED:           // 地址解析成功事件    
            ret = addr_handler(cma_id->context);
            break;
        case RDMA_CM_EVENT_MULTICAST_JOIN:          // 加入多播组成功事件
            ret = join_handler(cma_id->context, &event->param.ud);
            break;
        case RDMA_CM_EVENT_ADDR_ERROR:              // 各类错误事件
        case RDMA_CM_EVENT_ROUTE_ERROR:
        case RDMA_CM_EVENT_MULTICAST_ERROR:
            printf("mckey: event: %s, error: %d\n", 
                rdma_event_str(event->event), 
                event->status);
            connect_error();
            ret = event->status;
            break;
        case RDMA_CM_EVENT_DEVICE_REMOVAL:
            /* Cleanup will occur after test completes. */
            break;
        default:
            break;
    }
    return ret;
}




static void destroy_node(struct cmatest_node *node) {
    if (!node->cma_id)
        return;

    if (node->ah)
        ibv_destroy_ah(node->ah);

    if (node->cma_id->qp)
        rdma_destroy_qp(node->cma_id);

    if (node->cq)
        ibv_destroy_cq(node->cq);

    if (node->mem) {
        ibv_dereg_mr(node->mr);
        free(node->mem);
    }

    if (node->pd)
        ibv_dealloc_pd(node->pd);
    /* Destroy the RDMA ID after all device resources */
    rdma_destroy_id(node->cma_id);
}




static int alloc_nodes(void) {
    int ret, i;
    test.nodes = malloc(sizeof *test.nodes * connections);

    if (!test.nodes) {
        printf("mckey: unable to allocate memory for test nodes\n");
        return -ENOMEM;
    }

    memset(test.nodes, 0, sizeof *test.nodes * connections);
    for (i = 0; i < connections; i++) {
        test.nodes[i].id = i;
        ret = rdma_create_id(test.channel, &test.nodes[i].cma_id, &test.nodes[i], port_space);
        if (ret)
            goto err;
    }
    return 0;

err:
    while (--i >= 0)
        rdma_destroy_id(test.nodes[i].cma_id);
    free(test.nodes);

    return ret;
}




static void destroy_nodes(void) {
    int i;
    for (i = 0; i < connections; i++)
        destroy_node(&test.nodes[i]);
    free(test.nodes);
}



/**
 *! @brief 轮询完成队列 (CQ) 以获取工作完成 (WC) 通知。
 */
static int poll_cqs(void) {
    struct ibv_wc wc[8];
    int done, i, ret;

    for (i = 0; i < connections; i++) {
        if (!test.nodes[i].connected)
            continue;

        // 循环 poll CQ，直到收到 message_count 个完成事件。
        for (done = 0; done < message_count; done += ret) {
            // ibv_poll_cq 从 CQ 中取出已完成的 WC。
            // 这是一个非阻塞调用，如果没有完成事件，它会立即返回 0。
            ret = ibv_poll_cq(test.nodes[i].cq, 8, wc);

            if (ret < 0) {
                printf("mckey: failed polling CQ: %d\n", ret);
                return ret;
            }
        }
    }
    return 0;
}



/**
 *! @brief 主事件循环。从事件通道中获取并处理 CM 事件。
 */
static int connect_events(void) {
    struct rdma_cm_event *event;
    int ret = 0;

    // 当还有连接任务未完成且没有发生错误时，循环
    while (test.connects_left && !ret) {
        // rdma_get_cm_event 是一个阻塞调用，等待下一个 CM 事件的到来。
        ret = rdma_get_cm_event(test.channel, &event);

        if (!ret) {
            // 调用主处理函数
            ret = cma_handler(event->id, event);

            // 确认事件已被处理，释放事件资源。
            rdma_ack_cm_event(event);
        }
    }
    return ret;
}




static int get_addr(char *dst, struct sockaddr *addr) {
    struct addrinfo *res;
    int ret;
    ret = getaddrinfo(dst, NULL, NULL, &res);
    if (ret) {
        printf("getaddrinfo failed - invalid hostname or IP address\n");
        return ret;
    }
    memcpy(addr, res->ai_addr, res->ai_addrlen);
    freeaddrinfo(res);
    return ret;
}



// =============================================================================
//                             主逻辑
// =============================================================================
static int run(void) {
    int i, ret;
    printf("mckey: starting %s\n", is_sender ? "client" : "server");


    // 1. 解析字符串地址为 sockaddr 结构
    if (src_addr) {
        ret = get_addr(src_addr, (struct sockaddr *) &test.src_in);
        if (ret)
            return ret;
    }

    ret = get_addr(dst_addr, (struct sockaddr *) &test.dst_in);
    if (ret)
        return ret;

    
    printf("mckey: joining\n");
    for (i = 0; i < connections; i++) {
        // (可选) 绑定本地地址
        if (src_addr) {
            ret = rdma_bind_addr(test.nodes[i].cma_id, test.src_addr);
            
            if (ret) {
                printf("mckey: addr bind failure: %d\n", ret);
                connect_error();
                return ret;
            }
        }

        if (unmapped_addr)
            ret = addr_handler(&test.nodes[i]);
        else
            // 2. **关键**: 调用 rdma_resolve_addr 开始地址解析过程。
            //    这是一个异步操作，成功后会产生一个 RDMA_CM_EVENT_ADDR_RESOLVED 事件。
            //    参数：
            //      - cma_id: 要解析的 CM ID。
            //      - src_addr: 本地源地址 (可选)。
            //      - dst_addr: 目标多播地址。
            //      - timeout_ms: 超时时间。
            ret = rdma_resolve_addr(test.nodes[i].cma_id, 
                    test.src_addr, test.dst_addr, 2000);

        if (ret){
            printf("mckey: resolve addr failure: %d\n", ret);
            connect_error();
            return ret;
        }
    }

    // 3. 进入事件处理循环，等待所有节点都加入多播组
    ret = connect_events();
    if (ret)
        goto out;

    // **重要**: 等待一段时间，给 InfiniBand 交换机的子网管理器(SM)足够的时间
    // 来配置交换机的多播转发表。如果没有这个延迟，早期的多播包可能会丢失。
    sleep(3);
    
    // 4. 开始数据传输
    if (message_count) {
        if (is_sender) {
            printf("initiating data transfers\n");
        
            // 发送端投递发送请求
            for (i = 0; i < connections; i++) {
                // 注意：这里 post_sends 的 remote_qpn 实际上会被多播的 AH 覆盖，
                // 真正的目标 QPN 是 0xFFFFFF。
                ret = post_sends(&test.nodes[i], 0);
                
                if (ret)
                    goto out;
            }
        } else {
            printf("receiving data transfers\n");
            // 接收端轮询 CQ 等待数据到达
            ret = poll_cqs();
            if (ret)
                goto out;
        }
        printf("data transfers complete\n");
    }
out:
    // 5. 离开多播组
    for (i = 0; i < connections; i++) {
        ret = rdma_leave_multicast(test.nodes[i].cma_id, test.dst_addr);
        
        if (ret)
            printf("mckey: failure leaving: %d\n", ret);
    }
    return ret;
}




// =============================================================================
//                             程序入口
// =============================================================================
int main(int argc, char **argv)
{
    int op, ret;

    // 1. 使用 getopt 解析命令行参数
    while ((op = getopt(argc, argv, "m:M:sb:c:C:S:p:")) != -1) {
        switch (op) {
            case 'm':                // 多播地址
                dst_addr = optarg;
                break;
            case 'M':
                unmapped_addr = 1;
                dst_addr = optarg;
                break;
            case 's':                // 设为发送端
                is_sender = 1;
                break;
            case 'b':
                src_addr = optarg;
                test.src_addr = (struct sockaddr *) &test.src_in;
                break;
            case 'c':
                connections = atoi(optarg);
                break;
            case 'C':               // 消息数量
                message_count = atoi(optarg);
                break;
            case 'S':               // 消息大小
                message_size = atoi(optarg);
                break;
            case 'p':
                port_space = strtol(optarg, NULL, 0);
                break;
            default:
            printf("usage: %s\n", argv[0]);
            printf("\t-m multicast_address\n");
            printf("\t[-M unmapped_multicast_address]\n"
            "\t replaces -m and requires -b\n");
            printf("\t[-s(ender)]\n");
            printf("\t[-b bind_address]\n");
            printf("\t[-c connections]\n");
            printf("\t[-C message_count]\n");
            printf("\t[-S message_size]\n");
            printf("\t[-p port_space - %#x for UDP (default), %#x for IPOIB]\n", 
                    RDMA_PS_UDP, 
                    RDMA_PS_IPOIB);
            exit(1);
        }
    }


    // 2. 初始化全局测试结构体
    test.dst_addr = (struct sockaddr *) &test.dst_in;
    test.connects_left = connections;

    // 3. 创建 RDMA 事件通道
    test.channel = rdma_create_event_channel();

    if (!test.channel) {
        printf("failed to create event channel\n");
        exit(1);
    }

    // 4. 分配节点资源
    if (alloc_nodes())
        exit(1);


    // 5. 执行主逻辑
    ret = run();
    

    // 6. 清理资源并退出
    printf("test complete\n");
    destroy_nodes();
    rdma_destroy_event_channel(test.channel);
    printf("return status %d\n", ret);
    return ret;
}